import json
from typing import List, Dict
from question.chat_chain.config import CFG
from question.llm_service.qwenService import QwenService
def clean_llm_response(response: str) -> str:
    """
    清理LLM返回结果中的格式标记（如```json前缀和```后缀）

    Args:
        response: LLM返回的原始字符串

    Returns:
        清理后的字符串
    """
    # 去除首尾空白字符
    cleaned = response.strip()

    # 移除开头的```json标记
    if cleaned.startswith('```json'):
        cleaned = cleaned[len('```json'):].strip()

    # 移除结尾的```标记
    if cleaned.endswith('```'):
        cleaned = cleaned[:-len('```')].strip()

    return cleaned
def handle_ie_llm(llm_ie):
    """
    处理大模型抽取的实体和关系数据，存入Neo4j数据库
    :param llm_ie: LLMIEUtil.predict()返回的结果，格式为{"entities": [], "relations": []}
    """
    print("LLM数据处理完成，开始入库！")
    syz = set()  # 存储(主体, 主体标签, 关系, 客体, 客体标签)
    entities = set()  # 存储(实体值, 实体标签)

    # 提取实体
    for entity in llm_ie.get("entities", []):
        node_data = entity.get("node_data")
        node_label = entity.get("node_label")
        if node_data and node_label:  # 过滤无效实体
            entities.add((node_data, node_label))

    # 提取关系及关联实体
    for relation in llm_ie.get("relations", []):
        start_node_data = relation.get("start_node_data")
        start_label = relation.get("start_label")
        rel_type = relation.get("relation_type")
        end_node_data = relation.get("end_node_data")
        end_label = relation.get("end_label")

        # 过滤无效关系
        if not all([start_node_data, start_label, rel_type, end_node_data, end_label]):
            continue

        # 确保关系两端的实体被添加到实体集合
        entities.add((start_node_data, start_label))
        entities.add((end_node_data, end_label))
        # 添加关系到集合
        syz.add((start_node_data, start_label, rel_type, end_node_data, end_label))

    # 批量插入实体节点到主从数据库
    for node_data, node_label in entities:
        # 处理节点名称中的单引号，避免Cypher语法错误
        node_data_escaped = node_data.replace("'", "\\'")
        cypher = f"MERGE (a:{node_label} {{name: '{node_data_escaped}'}})"
        print(cypher)
        # g_master.run(cypher)
        # g_from.run(cypher)

    # 批量插入关系到主从数据库
    for start_data, start_label, rel_type, end_data, end_label in syz:
        # 处理特殊字符转义
        start_data_escaped = start_data.replace("'", "\\'")
        end_data_escaped = end_data.replace("'", "\\'")
        cypher = f"""
            MERGE (a:{start_label} {{name: '{start_data_escaped}'}})
            MERGE (b:{end_label} {{name: '{end_data_escaped}'}})
            MERGE (a)-[r:{rel_type}]->(b)
        """
        # g_master.run(cypher)
        # g_from.run(cypher)
        print(cypher)

    print("LLM抽取结果入库成功")
class LLMIEUtil:
    def __init__(self):
        self.config = CFG()
        self.ner_schema = ['设备', '部件', '检查内容', '故障现象', '故障原因', '排故流程', '维修方案']
        self.rel_schema = [
            "设备-包含-部件",
            "部件-包括-检查内容",
            "部件-产生-故障现象",
            "检查内容-涉及-故障现象",
            "故障现象-排故-排故流程",
            "故障原因-导致-故障现象",
            "故障原因-相关-部件",
            "故障原因-建议-维修方案",
            "排故流程-对应-故障原因"
        ]
        self.input_example = "最新型号的皮带机，部件包括螺丝、垫片等"
        # 修正实体示例：移除外层字典，直接使用数组（与Format要求一致）
        self.entities_exeample = [
            {
                "node_label": "设备",
                "node_data": "皮带机"
            },
            {
                "node_label": "部件",
                "node_data": "螺丝"
            },
            {
                "node_label": "部件",
                "node_data": "垫片"
            }
        ]
        # 修正关系示例：移除外层字典，直接使用数组（与Format要求一致）
        self.relations_exeample = [
            {
                "start_label": "设备",
                "start_node_data": "皮带机",
                "end_label": "部件",
                "end_node_data": "螺丝",
                "relation_type": "包含"  # 明确包含relation_type
            },
            {
                "start_label": "设备",
                "start_node_data": "皮带机",
                "end_label": "部件",
                "end_node_data": "垫片",
                "relation_type": "包含"
            }
        ]
        # self.llm = QwenService()
        # self.llm.load_model(model_name_or_path="D:\code\model\llm\Qwen2.5-1.5B-Instruct", isqiamtize4=False,
        #                     device="cuda")

        self.llm = QwenService()
        # 仅在未加载模型时才加载（避免重复加载）
        # self.llm.load_model(model_name_or_path=self.config.llm_model_name
        #                                 , isqiamtize4=False, device=self.config.device)
    def extract_entities(self, text: str) -> List[Dict[str, str]]:
        prompt = (
            "<Task>你是一名自然语言处理专家，擅长实体识别任务。"
            "给定一段文本和实体类型列表，请抽取文本中存在的实体。</Task>\n"
            f"<Input>{text}</Input>\n"
            f"<Entity Types>{self.ner_schema}</Entity Types>\n"
            "<Example>只参考下面例子的格式不参考内容\n"
            #f"  <Input Example>{self.input_example}</Input Example>\n"
            f"  <Output Example>{self.entities_exeample}</Output Example>\n"
            "</Example>\n"
            "<Format>请以JSON数组格式返回结果，仅包含实体对象。每个实体必须包含键：node_label（实体类型）、node_data（实体值）</Format>"
        )
        response = self.llm._call(prompt)
        response = clean_llm_response(response)
        print("实体：", response)

        # 解析优化：使用json模块，处理单引号转双引号
        try:
            # 替换单引号为双引号，修复JSON格式
            response = response.replace("'", '"')
            entities = json.loads(response)

            if isinstance(entities, list):
                val_entities = []
                for e in entities:
                    if isinstance(e, dict) and "node_label" in e and "node_data" in e:
                        val_entities.append(e)
                return val_entities
        except json.JSONDecodeError as e:
            print(f"实体JSON解析错误: {e}")
        except Exception as e:
            print(f"实体处理错误: {e}")
        return []

    def extract_relations(self, text: str, entities: List[Dict[str, str]]) -> List[Dict[str, str]]:
        entities_str = json.dumps(entities, ensure_ascii=False)
        prompt = (
            "<Task>你是一名自然语言处理专家，擅长关系抽取任务。"
            "给定文本、实体列表和关系类型列表，抽取文本中存在的关系。</Task>\n"
            f"<Input>{text}</Input>\n"
            f"<Entities>{entities_str}</Entities>\n"
            f"<Relation Types>{self.rel_schema}</Relation Types>\n"
            "<Example>只参考下面例子的格式不参考内容\n"
            #f"  <Input Example>{self.input_example}</Input Example>\n"
            f"  <Output Example>{self.relations_exeample}</Output Example>\n"
            "</Example>\n"
            "<Format>请以JSON数组格式返回结果，仅包含关系对象。每个关系必须包含键：start_label（起始实体类型）、start_node_data（起始实体值）、end_label（结束实体类型）、end_node_data（结束实体值）、relation_type（关系类型）</Format>"
        )
        response = self.llm._call(prompt)
        response = clean_llm_response(response)
        print("关系：", response)

        # 解析优化：使用json模块
        try:
            response = response.replace("'", '"')
            relations = json.loads(response)

            if isinstance(relations, list):
                valid_relations = []
                for r in relations:
                    required_keys = ["start_label", "start_node_data", "end_label", "end_node_data", "relation_type"]
                    if isinstance(r, dict) and all(k in r for k in required_keys):
                        rel_str = f"{r['start_label']}-{r['relation_type']}-{r['end_label']}"
                        if rel_str in self.rel_schema:
                            valid_relations.append(r)
                return valid_relations
        except json.JSONDecodeError as e:
            print(f"关系JSON解析错误: {e}")
        except Exception as e:
            print(f"关系处理错误: {e}")
        return []
    def predict(self, text: str) -> Dict[str, List[Dict[str, str]]]:
        """完整抽取流程：先抽实体，再抽关系"""
        print("正在抽取：",text)
        entities = self.extract_entities(text)
        relations = self.extract_relations(text, entities)
        return {
            "entities": entities,
            "relations": relations
        }
if __name__ == '__main__':
    llmie = LLMIEUtil()
    sentence = "刮板机包含的设备部件有：机头部电动机，机头部减速器，机头部链轮，机头部液压系统，中间部销排座，中间部销排。"
    res=llmie.predict(sentence)
    print(res)
    handle_ie_llm(res)
